#ifdef PIPELINE

#ifndef __PIPELINE_H__
#define __PIPELINE_H__

#include <iostream>
#include <string>
#include <utility>
#include <functional>

#include <boost/thread.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/make_shared.hpp>
#include <boost/thread/mutex.hpp>

#include "caffe/net.hpp"
#include "caffe/internal_thread.hpp"
#include "caffe/util/blocking_queue.hpp"

namespace caffe {

  // enum class Direction {
  //                       FORWARD = 0,
  //                       BACKWARD = 1
  // };

  // typedef struct BatchBoundary {
  //   int start_idx_;
  //   int end_idx_;
  // } BatchBoundary_t;

  // Auxiliary function types for the Pipeline {FORWARD, BACKWARD} functions.
  using BackwardCallback_t = std::function<void(Net&, int, int, bool)>;
  using ForwardCallback_t = std::function<void(Net&, int, int)>;

  /**
   * Pipeline stage class. Used to implement a pipeline execution strategy
   */
  class PipeStage
  {
    friend class Pipeline;
  public:
    PipeStage(boost::shared_ptr<Net> net, const int base_idx, const int start_layer,
              const int end_layer, const int batch_size);
    PipeStage(boost::shared_ptr<Net> net, const int base_idx, const int start_layer,
              const int end_layer, const int batch_size, Direction direction);

    ~PipeStage();

    std::string StageString();

    bool is_started() const {
      return stage_thread_->joinable();
    }

    void Wait() {
      try {
        if (is_started()) {
          stage_thread_->join();
        }
      } catch (boost::thread_interrupted&) {
      } catch (std::exception& e) {
        LOG(FATAL) << "Thread exception: " << e.what();
      }
    }

    void StageStart(vector<float>* loss, bool main_thread = false);

  protected:
    void StageThreadEntry(vector<float>* loss);

    boost::shared_ptr<Net> net_;

    int base_idx_;
    int start_layer_idx_, end_layer_idx_;

    boost::shared_ptr<BlockingQueue<BatchBoundary_t>> stage_inputs_;
    boost::shared_ptr<BlockingQueue<BatchBoundary_t>> stage_outputs_;

    int batch_size_;

    Direction direction_;

    boost::shared_ptr<PipeStage> next_stage_;

  private:
    boost::shared_ptr<boost::thread> stage_thread_;
  }; // PipeStage class.

  /**
   * @brief Pipeline class.
   */
  class Pipeline
  {
  public:
    Pipeline(boost::shared_ptr<Net> net, Direction direction, int start_layer, int end_layer,
             int batch_size) :
      direction_(direction), net_(net), start_layer_idx_(start_layer), end_layer_idx_(end_layer),
      batch_size_(batch_size), macro_batch_size_(0) {
      pipe_inputs_ = boost::make_shared<BlockingQueue<BatchBoundary_t>>();
      pipe_outputs_ = boost::make_shared<BlockingQueue<BatchBoundary_t>>();

      this->layer_loss_ = std::vector<float>(end_layer - start_layer + 1, 0.0);

      Init();
    }

    ~Pipeline();

    int LayerNum() {
      return (this->end_layer_idx_ - this->start_layer_idx_ + 1);
    }
    std::string PipelineString();

    void InitPipeInputs();

    void ExecuteForward(ForwardCallback_t ForwardCallback, float *loss = nullptr) {
      InitPipeInputs();             // Init the pipeline inputs.
      for (auto i = 1; i < stages_.size(); ++ i) {
        stages_[i]->StageStart(&layer_loss_);
      }
      stages_[0]->StageStart(&layer_loss_, true);
      // for (auto &stage : stages_) {
      //   stage->StageStart(&layer_loss_);
      // }

      this->WaitAll();

      if (loss != nullptr) {
        for (auto layer_loss : layer_loss_) {
          *loss += layer_loss;
        }
      }

      ForwardCallback(*(this->net_), this->start_layer_idx_,
                      this->end_layer_idx_);
    }

    void ExecuteBackward(BackwardCallback_t BackwardCallback,
                         bool apply_update = false) {
      for (auto &stage : stages_) {
        stage->StageStart(nullptr);
      }

      this->WaitAll();
      BackwardCallback(*(this->net_), this->start_layer_idx_,
                       this->end_layer_idx_, apply_update);
    }

  protected:
    void WaitAll() {
      // for (auto &stage : stages_) {
      //   stage->Wait();
      // }
      for (auto i = 1; i < stages_.size(); ++ i) {
        stages_[i]->Wait();
      }
    }

    virtual void Init();
    virtual void InitStages();

    Direction direction_;
    boost::shared_ptr<Net> net_;
    std::vector<boost::shared_ptr<PipeStage>> stages_;

    int start_layer_idx_, end_layer_idx_;
    int batch_size_, macro_batch_size_;

    std::vector<float> layer_loss_;

    boost::shared_ptr<BlockingQueue<BatchBoundary_t>> pipe_inputs_;
    boost::shared_ptr<BlockingQueue<BatchBoundary_t>> pipe_outputs_;
  }; // Pipeline class.

} // namespace caffe.

#endif /* __PIPELINE_H__ */

#endif  // PIPELINE
